// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <optional>
#include <set>
#include <unordered_map>

#include <pollux/plan/reference_query_runner.h>

using namespace kumo::pollux::plan;

namespace kumo::pollux::exec::test {
    class DuckQueryRunner : public ReferenceQueryRunner {
    public:
        explicit DuckQueryRunner(memory::MemoryPool *aggregatePool);

        RunnerType runnerType() const override {
            return RunnerType::kDuckQueryRunner;
        }

        /// Skip Timestamp, Varbinary, Unknown, and IntervalDayTime types. DuckDB
        /// doesn't support nanosecond precision for timestamps or casting from Bigint
        /// to Interval.
        ///
        /// TODO Investigate mismatches reported when comparing Varbinary.
        const std::vector<TypePtr> &supportedScalarTypes() const override;

        const std::unordered_map<std::string, DataSpec> &
        aggregationFunctionDataSpecs() const override;

        /// Specify names of aggregate function to exclude from the list of supported
        /// functions. Used to exclude functions that are non-determonistic, have bugs
        /// or whose semantics differ from Pollux.
        void disableAggregateFunctions(const std::vector<std::string> &names);

        /// Supports AggregationNode and WindowNode with optional ProjectNode on top.
        /// Assumes that source of AggregationNode or Window Node is 'tmp' table.
        std::optional<std::string> toSql(const core::PlanNodePtr &plan) override;

        // Converts 'plan' into an SQL query and executes it. Result is returned as a
        // MaterializedRowMultiset with the ReferenceQueryErrorCode::kSuccess if
        // successful, or an std::nullopt with a ReferenceQueryErrorCode if the query
        // fails.
        std::pair<
            std::optional<std::multiset<std::vector<pollux::variant> > >,
            ReferenceQueryErrorCode>
        execute(const core::PlanNodePtr &plan) override;

    private:
        std::unordered_set<std::string> aggregateFunctionNames_;
    };
} // namespace kumo::pollux::exec::test
